#include <stdio.h>
#include <stdlib.h>
#include <pthread.h>
#include <time.h>
#include <sys/time.h>

#define MAX_THREADS 16

int m, n, k;
int **A, **B, **C;
int num_threads;

// 线程参数结构体
typedef struct {
    int thread_id;
} ThreadParam;

// 矩阵乘法线程函数
void *matrix_multiply(void *arg) {
    ThreadParam *param = (ThreadParam *)arg;
    int id = param->thread_id;
    
    // 计算每个线程负责的行范围
    int rows_per_thread = m / num_threads;
    int start_row = id * rows_per_thread;
    int end_row = (id == num_threads - 1) ? m : start_row + rows_per_thread;
    
    for (int i = start_row; i < end_row; i++) {
        for (int j = 0; j < k; j++) {
            C[i][j] = 0;
            for (int l = 0; l < n; l++) {
                C[i][j] += A[i][l] * B[l][j];
            }
        }
    }
    
    pthread_exit(NULL);
}

// 初始化矩阵
void init_matrix(int ***matrix, int rows, int cols) {
    *matrix = (int **)malloc(rows * sizeof(int *));
    for (int i = 0; i < rows; i++) {
        (*matrix)[i] = (int *)malloc(cols * sizeof(int));
        for (int j = 0; j < cols; j++) {
            (*matrix)[i][j] = rand() % 100;
        }
    }
}

// 打印矩阵
void print_matrix(int **matrix, int rows, int cols) {
    for (int i = 0; i < rows; i++) {
        for (int j = 0; j < cols; j++) {
            printf("%d ", matrix[i][j]);
        }
        printf("\n");
    }
}

int main(int argc, char *argv[]) {
    if (argc != 5) {
        printf("Usage: %s <m> <n> <k> <num_threads>\n", argv[0]);
        return 1;
    }
    
    m = atoi(argv[1]);
    n = atoi(argv[2]);
    k = atoi(argv[3]);
    num_threads = atoi(argv[4]);
    
    if (num_threads < 1 || num_threads > MAX_THREADS) {
        printf("Number of threads must be between 1 and %d\n", MAX_THREADS);
        return 1;
    }
    
    srand(time(NULL));
    
    // 初始化矩阵
    init_matrix(&A, m, n);
    init_matrix(&B, n, k);
    init_matrix(&C, m, k);
    
    // 创建线程
    pthread_t threads[MAX_THREADS];
    ThreadParam params[MAX_THREADS];
    
    struct timeval start, end;
    gettimeofday(&start, NULL);
    
    for (int i = 0; i < num_threads; i++) {
        params[i].thread_id = i;
        pthread_create(&threads[i], NULL, matrix_multiply, &params[i]);
    }
    
    // 等待所有线程完成
    for (int i = 0; i < num_threads; i++) {
        pthread_join(threads[i], NULL);
    }
    
    gettimeofday(&end, NULL);
    double time_taken = (end.tv_sec - start.tv_sec) * 1e6;
    time_taken = (time_taken + (end.tv_usec - start.tv_usec)) * 1e-6;
    
    // 输出结果
    printf("Matrix A:\n");
    print_matrix(A, m, n);
    
    printf("\nMatrix B:\n");
    print_matrix(B, n, k);
    
    printf("\nMatrix C (Result):\n");
    print_matrix(C, m, k);
    
    printf("\nTime taken: %f seconds\n", time_taken);
    
    // 释放内存
    for (int i = 0; i < m; i++) free(A[i]);
    for (int i = 0; i < n; i++) free(B[i]);
    for (int i = 0; i < m; i++) free(C[i]);
    free(A);
    free(B);
    free(C);
    
    return 0;
}